#!/bin/bash

# Copyright 2022 Balena Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

#
# OS migration
#
# Expects to find flash-boot and a mounted flash-rootA rootfs in the boot disk

# shellcheck disable=SC1091
. /usr/libexec/os-helpers-logging
# shellcheck disable=SC1091
. /usr/libexec/os-helpers-fs
# shellcheck disable=SC1091
. /usr/sbin/balena-config-defaults

memory_check() {
    _size="${1}"
    sync && echo 3 > /proc/sys/vm/drop_caches
    swapoff -a && swapon -a
    _memfree=$(awk '/MemFree/{free=$2} END{print (free*1024)}' /proc/meminfo)
    if [ "${_memfree}" -lt "${_size}" ]; then
        fail "Not enough memory: requested ${_size}, available ${_memfree}"
    fi
    return 0
}

# Enable module function
migrate_enabled() {
    # shellcheck disable=SC2154
    if [ "$bootparam_flasher" = "true" ]; then
        # flash-rootA has been identified as rootfs

        _flash_file_name=$(basename "${FLASHER_FILEFLAG}")
        if [ -f "${ROOTFS_DIR}/flash-boot/${_flash_file_name}" ]; then
            # shellcheck disable=SC1090
            . "${ROOTFS_DIR}/etc/resin-init-flasher.conf"

            # Migration use cases are:
            #
            # * Booting with only one disk - migrate if this disk is the one
            #   configured to program into
            # * Booting with multiple disks - only migrate if explicitely
            #   configured to do so as the user might want to install on
            #   alternative disks
            #
            internal_dev=$(get_internal_device "${INTERNAL_DEVICE_KERNEL}")
            boot_dev=$(findmnt --noheadings --canonicalize --output SOURCE "${ROOTFS_DIR}" | xargs lsblk -no pkname)
            FLASH_BOOT_DEVICE=$(get_dev_path_in_device_with_label "/dev/${boot_dev}" "flash-boot")
            if [ -n "${FLASH_BOOT_DEVICE}" ]; then
                FLASH_BOOT_MOUNT="/tmp/flash-boot"
                mkdir -p "${FLASH_BOOT_MOUNT}"
                mount "${FLASH_BOOT_DEVICE}" "${FLASH_BOOT_MOUNT}"
                if [ "$bootparam_migrate" = "true" ] || jq -re '.installer.migrate.force' "${FLASH_BOOT_MOUNT}/config.json" > /dev/null; then
                    _migrate=1
                    info "Migration requested in configuration"
                elif jq -re '.installer.secureboot' "${FLASH_BOOT_MOUNT}/config.json" > /dev/null; then
                    _migrate=1
                    info "Migration forced because secure boot is enabled"
                fi

                if [ "${_migrate}" = "1" ] && target_devices=$(jq -re '.installer.target_devices' "${FLASH_BOOT_MOUNT}/config.json"); then
                    info "Configured target_devices: $target_devices"
                    INTERNAL_DEVICE_KERNEL="${target_devices}"
                    internal_dev=$(get_internal_device "${INTERNAL_DEVICE_KERNEL}")
                fi
            else
                fail "Flash boot partition not found in ${internal_dev}"
            fi

            # Migrate if configured to do so or if there is only one disk
            # (excluding RAM and swap devices)
            if [[ ( -n "${_migrate}" && "${_migrate}" = "1" ) || $(lsblk -nde 251,1 | wc -l) -eq "1" ]]; then
                info "Running migration on ${internal_dev}..."
                return 0
            fi
        fi

        if is_secured && [ "${_migrate}" != "1" ]; then
            fail "Locked devices can only be installed from memory"
        fi

        # Leave flashing after pivot-rooting
        if mountpoint "${FLASH_BOOT_MOUNT}"; then
            umount "${FLASH_BOOT_MOUNT}" > /dev/null || true
        fi
    fi
    # standard boot
    return 1
}

# Main module function
migrate_run() {
    # Find the raw image in the rootfs partition
    image=$(find "${ROOTFS_DIR}" -xdev -type f -name "${BALENA_IMAGE}")
    kernel_images=$(find "${ROOTFS_DIR}" -xdev -type f -name "@@KERNEL_IMAGETYPE@@*")
    if [ -n "${image}" ]; then
        EXTERNAL_DEVICE_BOOT_PART_MOUNTPOINT="${BALENA_BOOT_MOUNTPOINT}"
        if findmnt "${FLASH_BOOT_MOUNT}" > /dev/null; then
            _image_size=$(estimate_size_in_zram "${image}")
            _flash_boot_size=$(du -cb ${FLASH_BOOT_MOUNT} | awk '/total/{print $1}')
            # shellcheck disable=SC2086
            _kernel_images_size=$(du -cb ${kernel_images} | awk '/total/{print $1}')
            _total_size=$(("$_image_size" + "$_flash_boot_size" + "$_kernel_images_size"))
            memory_check "${_total_size}"
        else
            fail "Flash boot partition not found in ${internal_dev}"
        fi
        # Copy the raw image to memory
        cp "${image}" "/tmp"
        # If the image is signed, copy the signature to memory
        if [ -f "${image}.sig" ]; then
            cp "${image}.sig" "/tmp"
        fi

        # Copy the flasher boot partition into memory (contains configuration)
        mkdir -p "${EXTERNAL_DEVICE_BOOT_PART_MOUNTPOINT}"
        cp -r "${FLASH_BOOT_MOUNT}"/* "${EXTERNAL_DEVICE_BOOT_PART_MOUNTPOINT}/"

        # Copy the flasher kernel images to memory
        # shellcheck disable=SC2086
        cp -a ${kernel_images} "/tmp"

        # Need to source this again to set CONFIG_PATH correctly
        unset CONFIG_PATH
        # shellcheck disable=SC1091
        . /usr/sbin/balena-config-defaults

        mkdir -p "$(dirname "${CONFIG_PATH}")"
        _config_json_name=$(basename "${CONFIG_PATH}")
        cp "${FLASH_BOOT_MOUNT}/${_config_json_name}" "${CONFIG_PATH}"

        umount "${FLASH_BOOT_MOUNT}" > /dev/null || true
        sync "${EXTERNAL_DEVICE_BOOT_PART_MOUNTPOINT}"
        # Unmount the rootfs as we are going to program over
        umount "${ROOTFS_DIR}"

        # If booting from the same disk we want to program, reboot after flashing
        if [ "${internal_dev#/dev/}" = "${boot_dev}" ]; then
            if [ -f "/etc/resin-init-flasher.conf" ]; then
                echo "POSTINSTALL_REBOOT=1" >> "/etc/resin-init-flasher.conf"
            else
                fail "Flasher configuration not found"
            fi
        fi

        # Mount securityfs so that TPM event log is accessible to flasher
        mount -t securityfs securityfs /sys/kernel/security

        # Run flasher - should not return
        /usr/bin/resin-init-flasher
        # Something went wrong - kill the process to prevent exploits
        exit 1
        # Just being paranoid
        echo "c" > /proc/sysrq-trigger
        sleep infinity
    else
        # If recovery mode, wait for adbd to exit
        if [ -n "${ADBD_PID}" ]; then
            # adbd is not a child process so cannot wait
            while kill -0  "${ADBD_PID}" 2>/dev/null; do sleep 1; done
        fi
        fail "No ${BALENA_IMAGE} found in ${ROOTFS_DIR}"
    fi
}
